#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <windows.h>

#include "xdisp_gdi.h"

///小矩形块
#define GDI_SMALL_RECT_WIDTH    16
#define GDI_SMALL_RECT_HEIGHT   16

#define SAFE_RELEASE(a) if(a){ (a)->Release(); (a)=NULL; }
#define SAFE_FREE(a) if(a){ free(a); (a)=NULL; }

#define max(a,b)    ({ typeof(a) _a = (a); typeof(b) _b = (b); \
    _a > _b ? _a : _b; })


#define min(a,b)    ({ typeof(a) _a = (a); typeof(b) _b = (b); \
    _a < _b ? _a : _b; })

////
int log_printf(const char* pFormat, ...);


///GDI 抓屏
struct __gdi_cap_t
{
    int        cx;
    int        cy;
    int        line_bytes;
    int        line_stride;
    int        bitcount; ////
    HDC        memdc;
    HBITMAP    hbmp;
    ///
    byte*      buffer; ///
    byte*      back_buf; ///
};

struct __tbuf_t
{
    int    size;
    void*  buf;
};
struct __xdisp_t
{
    bool             quit;

    bool             is_pause_grab; //是否暂停抓屏
    int              grab_type; //// 抓屏方法， 0 自动选择， 1 mirror， 2 DX抓屏， 3 GDI抓屏

    __gdi_cap_t      gdi;

    /////
    LONG             sleep_msec;  ///

    ////
    HWND             hMessageWnd;
    HANDLE           h_thread;
    DWORD            id_thread;
    HANDLE           hEvt;

    ///
    DP_DISPLAYCHANGE display_change; ///
    DP_FRAME         frame;
    void*            param;

    ///
    __tbuf_t         t_arr[2];
    /////
};

static void tbuf_check(__tbuf_t* arr, int size )
{
    if (!arr->buf || arr->size < size) {
        void* b1 = malloc(size);
        if (arr->buf)free(arr->buf);
        arr->buf = b1;
        arr->size = size;
    }
}

static bool __init_gdi(__xdisp_t* dp, bool is_init)
{
    if (is_init) {
        ////
        DEVMODE devmode;
        memset(&devmode, 0, sizeof(DEVMODE));
        devmode.dmSize = sizeof(DEVMODE);
        devmode.dmDriverExtra = 0;
        BOOL bRet = EnumDisplaySettings(NULL, ENUM_CURRENT_SETTINGS, &devmode);
        if (!bRet)return false;
        dp->gdi.cx = devmode.dmPelsWidth;
        dp->gdi.cy = devmode.dmPelsHeight;
        dp->gdi.bitcount = devmode.dmBitsPerPel;
        if (dp->gdi.bitcount != 8 && dp->gdi.bitcount != 16 && dp->gdi.bitcount != 32 && dp->gdi.bitcount !=24)return false;
        dp->gdi.line_bytes = dp->gdi.cx * dp->gdi.bitcount / 8;
        dp->gdi.line_stride = (dp->gdi.line_bytes + 3) / 4 * 4;

        BITMAPINFOHEADER bi; memset(&bi, 0, sizeof(bi));
        bi.biSize = sizeof(bi);
        bi.biWidth = dp->gdi.cx;
        bi.biHeight = -dp->gdi.cy; //从上朝下扫描
        bi.biPlanes = 1;
        bi.biBitCount = dp->gdi.bitcount; //RGB
        bi.biCompression = BI_RGB;
        bi.biSizeImage = 0;
        BYTE bb[sizeof(BITMAPINFOHEADER) + sizeof(RGBQUAD) * 256];
        memcpy(bb, &bi, sizeof(bi));
        HDC hdc = GetDC(NULL); //屏幕DC
        if (dp->gdi.bitcount == 8) {//系统调色板
            PALETTEENTRY pal[256]; memset(pal, 0, sizeof(pal));///
            GetSystemPaletteEntries(hdc, 0, 256, pal); ///获取系统调色板
            RGBQUAD* PalColor = (RGBQUAD*)(bb + sizeof(bi));
            for (int i = 0; i < 256; ++i) {
                PalColor[i].rgbReserved = 0;
                PalColor[i].rgbRed = pal[i].peRed;
                PalColor[i].rgbGreen = pal[i].peGreen;
                PalColor[i].rgbBlue = pal[i].peBlue;
            }
        }

        dp->gdi.memdc = CreateCompatibleDC(hdc);
        dp->gdi.buffer = NULL;
        dp->gdi.hbmp = CreateDIBSection(hdc, (BITMAPINFO*)bb, DIB_RGB_COLORS, (void**)&dp->gdi.buffer, NULL, 0);
        ReleaseDC(NULL, hdc);
        if (!dp->gdi.buffer) {
            DeleteDC(dp->gdi.memdc); dp->gdi.memdc = 0;
            return false;
        }
        SelectObject(dp->gdi.memdc, dp->gdi.hbmp); ///
        dp->gdi.back_buf = (byte*)malloc(dp->gdi.line_stride* dp->gdi.cy);
        memset(dp->gdi.back_buf, 0, dp->gdi.line_stride*dp->gdi.cy);
        return true;
    }
    else {
        dp->gdi.buffer = 0;
        DeleteDC(dp->gdi.memdc); dp->gdi.memdc = NULL;
        DeleteObject(dp->gdi.hbmp); dp->gdi.hbmp = NULL;

        SAFE_FREE(dp->gdi.back_buf);
    }

    return false;
}



LRESULT CALLBACK xDispWindowProc(HWND hwnd, UINT uMsg, WPARAM  wParam, LPARAM lParam)
{
    ////
    switch (uMsg)
    {
    case WM_DISPLAYCHANGE:
    {
        __xdisp_t* dp = (__xdisp_t*)GetPropA(hwnd, "xDispData");
        if (dp) {
            ////先调用回调，
            dp->display_change(LOWORD(lParam), HIWORD(lParam), wParam, dp->param); /// callback

            if (dp->grab_type == GRAB_TYPE_GDI) {
                __init_gdi(dp, false);
                __init_gdi(dp, true);
            }
            /////
        }
        ///
        //		printf("--- WM_DISPLAYCHANGE\n");
    }
        break;

    case WM_DESTROY:
        PostQuitMessage(0);
        break;

    }
    return DefWindowProc(hwnd, uMsg, wParam, lParam);
}
static HWND createMsgWnd(__xdisp_t* dp)
{
    HMODULE hmod = GetModuleHandle(NULL);

    const char* cls_name = "xdisp_window_class";
    WNDCLASS wndClass = { 0,
                          xDispWindowProc,
                          0, 0,
                          hmod, ///
                          NULL,NULL,NULL,NULL,
                          cls_name
                        };
    RegisterClass(&wndClass);
    HWND hwnd = CreateWindowA(cls_name, "", WS_OVERLAPPEDWINDOW, 0, 0, 1, 1, NULL, NULL, hmod, NULL);
    dp->hMessageWnd = hwnd;
    if (!hwnd) {
        printf("NULL Message Window.\n");
    }
    SetProp(hwnd, "xDispData", dp);

    ////
    return hwnd;
}

///
static void __forceinline combine_rect_to_rgn(HRGN& region, int x1, int y1, int x2, int y2 )
{
    if (!region) {
        region = CreateRectRgn(x1,y1,x2,y2); ////
    }
    else {
        HRGN n_rgn = CreateRectRgn(x1, y1, x2, y2);
        if (!n_rgn) {
            printf("CreateRectRgn [%d, %d, %d, %d} err=%d\n", x1,y1,x2,y2, GetLastError() );
            return;
        }
        if (CombineRgn(region, region, n_rgn, RGN_OR) == NULLREGION) { //矩形框可能有重叠，因此首先组成HRGN区域，然后再计算出合并之后的矩形框
            printf("CombineRgn [%d, %d, %d, %d} err=%d\n", x1, y1, x2, y2, GetLastError());
            DeleteObject(region);
            region = NULL;
        }
        DeleteObject(n_rgn); ////
        /////
    }
}
static int combine_rectangle(byte* primary_buffer, byte* second_buffer, 
                             int line_stride, LPRECT rect, LPSIZE screen, int bitcount,
                             int SMALL_WIDTH, int SMALL_HEIGHT , HRGN & region )
{
    int x, y;
    int num_bit = (bitcount >> 3);
    byte* dst_buf = primary_buffer;
    byte* bak_buf = second_buffer;

    for (y = max(rect->top, 0); y < rect->bottom; y += SMALL_HEIGHT ) {
        int h = min(SMALL_HEIGHT, rect->bottom - y); //printf("h=%d, [ %d]\n",h, rect->bottom - y);
        for (x = max(rect->left, 0); x < rect->right; x += SMALL_WIDTH) {
            int w = min(SMALL_WIDTH, rect->right - x);
            int x2 = x + w;
            int y2 = y + h;
            if (x < 0) { x = 0; } else if (x >= screen->cx) { x = screen->cx; }
            if (y < 0) { y = 0; } else if (y >= screen->cy) { y = screen->cy; }
            if (x2 >= screen->cx) x2 = screen->cx; if (x2 <= x)continue;
            if (y2 >= screen->cy) y2 = screen->cy; if (y2 <= y)continue;

            w = x2 - x;
            h = y2 - y;
            int base_pos = line_stride*y + x*num_bit;
            int line_bytes = w*num_bit;

            int i;
            bool is_dirty = false; //printf("*** [%d, %d, %d, %d] -> {%d, %d, %d, %d}\n", rect->left,rect->top,rect->right,rect->bottom, x,y,x2,y2);
#define    REP_CD(II) \
    if( !is_dirty){ \
    for (i = II; i < h; i += 5) { /*隔行扫描*/ \
    int pos = base_pos + line_stride*i;\
    if (memcmp(dst_buf + pos, bak_buf + pos, line_bytes) != 0) { is_dirty = true; break; }\
        }\
        }
            REP_CD(0); REP_CD(1);
            REP_CD(2); REP_CD(3); REP_CD(4);

            if (is_dirty) {
                //printf(" dirty [%d, %d, %d, %d]\n", x,y,x2,y2);
                combine_rect_to_rgn(region, x, y, x2, y2); //printf(" --dirty [%d, %d, %d, %d]\n", x, y, x2, y2);////
            }
            else {

            }
            ////////
        }
    }

    return 0;
}
static int region_to_rectangle(HRGN region, byte* dst_buf, int line_stride, 
                               int cx, int cy, int bitcount, dp_rect_t** p_rc_array, int* p_rc_count , __tbuf_t* tarr)
{
    dp_rect_t* rc_array = NULL;
    int rc_count = 0;
    int num_bit = (bitcount >> 3);
    ///
    DWORD size = GetRegionData(region, NULL, NULL);
    tbuf_check(&tarr[0], size);

    LPRGNDATA rgnData = (LPRGNDATA)tarr[0].buf;
    if (GetRegionData(region, size, rgnData)) { ///
        DWORD cnt = rgnData->rdh.nCount; //printf("new cnt=%d\n", cnt);
        DWORD sz = cnt * sizeof(dp_rect_t);
        tbuf_check(&tarr[1], sz);

        dp_rect_t* rc = (dp_rect_t*)tarr[1].buf;
        //	printf("old cnt=%d, newcnt=%d\n", old_rc_cnt, cnt );
        ////
        LPRECT src_rc = (LPRECT)rgnData->Buffer;
        DWORD idx = 0; int i;
        for (i = 0; i < (int)cnt; ++i) {
            dp_rect_t* c = &rc[idx];;

            ///裁剪矩形框，保证在可视区域
            int x = src_rc[i].left;     if (x < 0) { x = 0; } else if (x >= cx) { x = cx; }
            int y = src_rc[i].top;      if (y < 0) { y = 0; } else if (y >= cy) { y = cy; }
            int x2 = src_rc[i].right;   if (x2 >= cx) x2 = cx; if (x2 <= x)continue;
            int y2 = src_rc[i].bottom;  if (y2 >= cy) y2 = cy; if (y2 <= y)continue;
            //printf("--- {%d, %d, %d, %d}\n", x, y, x2, y2);
            int w = x2 - x; //
            int h = y2 - y;
            c->rc.left = x; c->rc.top = y; c->rc.right = x2; c->rc.bottom = y2; ///

            c->line_buffer = (char*)dst_buf + line_stride*y + x*num_bit;  ///
            c->line_bytes = w * num_bit;
            c->line_nextpos = line_stride; /////
            c->line_count = h; ///
            ///
            ++idx; ///
            //	if(x2 > dp->buffer->cy)printf("--** right=%d,bottom=%d, cx=%d,cy=%d\n", c->rc.right,c->rc.bottom, dp->buffer->cx, dp->buffer->cy);
        }
        rc_array = rc;
        rc_count = idx; //
        //////
    }
    *p_rc_array = rc_array;
    *p_rc_count = rc_count;
    /////
    return 0;
}
/////





static void capture_gdi(__xdisp_t* dp)
{
    if (!dp->gdi.buffer || !dp->gdi.hbmp )return;
    ///
    HDC hdc = GetDC(NULL);
    BitBlt(dp->gdi.memdc, 0, 0, dp->gdi.cx, dp->gdi.cy, hdc, 0, 0, SRCCOPY | CAPTUREBLT);
    ReleaseDC(NULL, hdc);
    ///
    int line_bytes = dp->gdi.line_bytes;
    int line_stride = dp->gdi.line_stride;
    int cx = dp->gdi.cx; int cy = dp->gdi.cy;
    HRGN region = NULL;

    SIZE screenSize; screenSize.cx = dp->gdi.cx; screenSize.cy = dp->gdi.cy;
    RECT screeRc = {0,0, dp->gdi.cx, dp->gdi.cy };
    combine_rectangle( dp->gdi.buffer, dp->gdi.back_buf, dp->gdi.line_stride, &screeRc, &screenSize, dp->gdi.bitcount,
                       GDI_SMALL_RECT_WIDTH, GDI_SMALL_RECT_HEIGHT, region);


    ///
    int rc_count = 0;
    dp_rect_t* rc_array = NULL;
    if (region) {
        ////
        region_to_rectangle(region, dp->gdi.buffer, dp->gdi.line_stride, dp->gdi.cx, dp->gdi.cy,
                            dp->gdi.bitcount, &rc_array, &rc_count, dp->t_arr);

        DeleteObject(region);
        //////
    }

    if (rc_count > 0) {
        ///
        memcpy(dp->gdi.back_buf, dp->gdi.buffer, dp->gdi.line_stride*cy); //// copy to backbuffer
    }

    /////
    dp_frame_t frm;
    frm.rc_array = rc_array;
    frm.rc_count = rc_count;
    frm.cx = dp->gdi.cx;
    frm.cy = dp->gdi.cy;
    frm.line_bytes = line_bytes;
    frm.line_stride = line_stride;
    frm.bitcount = dp->gdi.bitcount;
    frm.buffer = (char*)dp->gdi.buffer;
    frm.length = line_stride*dp->gdi.cy; ///
    frm.param = dp->param;

    dp->frame(&frm); /// callback

}


static DWORD CALLBACK __loop_msg(void* _p)
{
    __xdisp_t* dp = (__xdisp_t*)_p;  /////

    dp->id_thread = GetCurrentThreadId(); ///

    ::CoInitialize(0);

    //////
    int grab_type = dp->grab_type;
    bool is_ok = false;
    is_ok = __init_gdi(dp, true);
    if (is_ok) {
        grab_type = GRAB_TYPE_GDI;
        dp->grab_type = grab_type; ////
    }

    if (grab_type == GRAB_TYPE_GDI) {
        is_ok = __init_gdi(dp, true); /////
    }
    ///
    if (!is_ok ) {
        printf("*** can not init capture screen type=%d\n", grab_type );
        dp->hMessageWnd = NULL; ///
        SetEvent(dp->hEvt);
        ::CoUninitialize();
        return 0;
    }

    ////

    HWND hwnd = createMsgWnd(dp); ///
    HANDLE hEvt = dp->hEvt; ///

    if (!hwnd ) {
        /////
        switch (grab_type) {

        case GRAB_TYPE_GDI:     __init_gdi(dp, false);     break;
        }
        ////
        if (hwnd)DestroyWindow(hwnd);
        dp->hMessageWnd = NULL; ///
        SetEvent(hEvt); ////
        CoUninitialize();
        return 0; ///
    }
    ////
    SetEvent(hEvt); ////

    ///
    long sleep_msec = dp->sleep_msec;

    MSG msg;
    LARGE_INTEGER counter; QueryPerformanceFrequency(&counter);
    LARGE_INTEGER tmp_s;   QueryPerformanceCounter(&tmp_s);
    while ( !dp->quit ) {
        ///
        BOOL bQuit = FALSE;
        while (PeekMessage(&msg, NULL, 0, 0, PM_REMOVE)) {
            if (msg.message == WM_QUIT) {
                bQuit = TRUE;
                break;
            }
            else {
                DispatchMessage(&msg); ///
            }
            /////
        }
        if (bQuit)break;
        /////
        sleep_msec = dp->sleep_msec;
        DWORD captureFrame_time = 0;

        if ( !dp->is_pause_grab  ) { ///

            LARGE_INTEGER frame_begin; QueryPerformanceCounter(&frame_begin);

            if (grab_type == GRAB_TYPE_GDI) {
                ///
                capture_gdi(dp); ///
            }

            /////
            LARGE_INTEGER frame_end; QueryPerformanceCounter(&frame_end);

            const int MIN_DT = 10; ///

            DWORD dt = (frame_end.QuadPart - frame_begin.QuadPart) * 1000 / counter.QuadPart;
            captureFrame_time = dt; ///
            if (dt >= dp->sleep_msec - MIN_DT) dt = dp->sleep_msec - MIN_DT; //printf("encode use %dms\n",dt);

            sleep_msec = dp->sleep_msec - dt;//重新计算下次需要等待的时间

            ////////
        }

        /////沉睡sleep_msec 毫秒
        LARGE_INTEGER s, e;
        QueryPerformanceCounter(&s);
        LARGE_INTEGER counter_delay;
        counter_delay.QuadPart = counter.QuadPart * sleep_msec / 1000;
        while ( !dp->quit ) {
            SleepEx(1,TRUE); ///
            QueryPerformanceCounter(&e);
            if (e.QuadPart - s.QuadPart >= counter_delay.QuadPart )break;
        }
        /////
        //		DWORD dt2 = (e.QuadPart - tmp_s.QuadPart) * 1000 / counter.QuadPart; tmp_s = e;  printf("$$$ TowFrame Interval [%d]ms , ProcessFrame [%d],nextSleep=[%d]\n", dt2, captureFrame_time, sleep_msec);
        ///////
    }

    ///
    CloseHandle(hEvt);

    /////
    switch (dp->grab_type) {

    case GRAB_TYPE_GDI:     __init_gdi(dp, false);     break;
    }

    ::CoUninitialize(); ///
    ///////
    return 0;
}

void* dp_create(dp_create_t* ct)
{
    __xdisp_t* dp = NULL;

    ////
    dp = new __xdisp_t;
    dp->quit = false;
    dp->is_pause_grab = false;
    dp->grab_type = ct->grab_type;


    memset(&dp->gdi, 0, sizeof(dp->gdi));
    ///
    dp->sleep_msec = 33; /// 30 fps

    dp->hMessageWnd = NULL;
    dp->hEvt = CreateEvent(NULL, TRUE, FALSE, NULL);

    dp->display_change = ct->display_change; ///
    dp->frame = ct->frame;
    dp->param = ct->param;

    memset(dp->t_arr, 0, sizeof(dp->t_arr));

    DWORD tid;
    dp->h_thread = CreateThread(NULL, 10*1024*1024, __loop_msg, dp, 0, &tid);

    if (!dp->h_thread) {
        CloseHandle(dp->hEvt);
        delete dp;
        //////
        return NULL;
    }
    ::SetThreadPriority(dp->h_thread, THREAD_PRIORITY_HIGHEST); //提升优先级

    ::WaitForSingleObject(dp->hEvt, INFINITE);
    ::ResetEvent(dp->hEvt); //
    if (!dp->hMessageWnd) {
        CloseHandle(dp->hEvt); ///
        CloseHandle(dp->h_thread); ///
        delete dp;
        ////
        return NULL;
    }

    //////
    return dp;
}

void dp_destroy(void* handle)
{
    __xdisp_t* dp = (__xdisp_t*)handle;
    if (!dp)return;
    //////
    dp->quit = true;
    SetEvent(dp->hEvt);
    DestroyWindow(dp->hMessageWnd); ////
    ::WaitForSingleObject(dp->h_thread, 8 * 1000);
    ::TerminateThread(dp->h_thread, 0);
    CloseHandle(dp->h_thread);
    CloseHandle(dp->hEvt);

    ////
    switch (dp->grab_type) {
    case GRAB_TYPE_GDI:     __init_gdi(dp, false);     break;
    }

    ////
    for (int i = 0; i < sizeof(dp->t_arr) / sizeof(dp->t_arr[0]); ++i) {
        if (dp->t_arr[i].buf)free(dp->t_arr[i].buf);
    }
    delete dp;
}

int dp_grab_interval(void* handle, int grab_msec )
{
    __xdisp_t* dp = (__xdisp_t*)handle;
    if (!dp)return -1;
    //////
    if (grab_msec < 10) grab_msec = 10;
    dp->sleep_msec = grab_msec;

    return 0;
}
int dp_grap_pause(void* handle, int is_pause)
{
    __xdisp_t* dp = (__xdisp_t*)handle;
    if (!dp)return -1;
    //////
    if (is_pause)dp->is_pause_grab = true;
    else dp->is_pause_grab = false;

    return 0;
}

//////////


#if 0

int dp_frame(dp_frame_t* frame)
{
    return 0;
}

int main(int argc, char** argv)
{
    dp_create_t ct;
    ct.is_mirror = 1;
    ct.frame = dp_frame; ///

    void* dp = dp_create(&ct);

    getchar();

    dp_destroy(dp); ///

    return 0;
}

#endif

